import os

import requests
from langchain_chroma import Chroma
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from langchain_deepseek import ChatDeepSeek
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_tavily import TavilySearch

import prompt_template


@tool
def empty_tools():
    """empty tools"""
    return "empty_tools"

@tool
def search_tools(query: str):
    """只有在需要实时信息或不知道的事情的时候才会使用这个工具。"""
    search_tool = TavilySearch(
        include_answer=True
    )
    result = search_tool.invoke(query)
    print(result)
    return result

def get_chromadb() -> Chroma:

    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': True}
    model_name = "BAAI/bge-m3"
    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )
    db = Chroma(
        persist_directory="chromadb",
        embedding_function=embeddings
    )
    return db

@tool
def rag_tools(query: str):
    """只有回答与2024年运势或者龙年运势相关的问题的时候，会使用这个工具."""

    db = get_chromadb()
    retriever = db.as_retriever(
        search_type="similarity_score_threshold",
        search_kwargs={
            "score_threshold": 0.3,
            "k": 1
        }
    )
    result = retriever.invoke(input=query)
    print(f"rag结果：{result}")
    return result

@tool
def bazi_cesuan(query: str):
    """只有做八字排盘时候才会使用这个工具，需要输入用户姓名、性别和出生年月日时，如果缺少用户姓名、性别和出生年月日时则不可用。"""

    prompt = ChatPromptTemplate.from_template(prompt_template.bazi_cesuan_args_prompt_template)
    llm = ChatDeepSeek(model="deepseek-chat", temperature=0, max_tokens=1024, verbose=True)
    chain = prompt | llm | JsonOutputParser()
    arg_result = chain.invoke({"query" : query})
    print(f"参数解析结果：{arg_result}")
    if arg_result is not None:
        arg_result["api_key"] = os.getenv("YUANXIANJU_API_KEY")
    url = os.getenv("YUANXIANJU_URL") + "/Bazi/cesuan"
    result = requests.post(url, data=arg_result)
    if result.status_code == 200:
        try:
            result = result.json()
            print(f"八字测算结果：{result}")
            return result
        except Exception as e:
            return "八字测算失败，可能是你忘记询问用户姓名、性别和出生年月日时了。"
    else:
        return "技术错误，请告诉用户稍后再试。"
